/**
 * 内存初始化的一些主要操作。
 */
#include <cds/string.h>
#include <cds/list.h>
#include <mm/physical_area.h>
#include <mm/pages.h>
#include <mm/early_alloc.h>
#include <mm/memory.h>

/// 原子操作需要重新实现。
#include <stdatomic.h>

/* 储存单个NUMA节点的内存信息。因为目前没有实现NUMA，所以其长度恒为1。 */
struct pglist_struct pglists[MAX_PGLISTS_NUM] = { 0 };
/* 记录有多少个 pglist 结构体。*/
static unsigned long pglists_count = 0;
/* 储存所有内存区域结构体的数组。其中的顺序可能与NUMA节点内的优先内存区域结构体数组不一致。 */
static struct zone_struct *zones = NULL;
/* 所有内存区域的数量。 */
static unsigned long zones_count = 0;


/**
 * @brief 初始化单个pglist结构体。
 */
static inline int __init_pglist_struct(
    struct pglist_struct *pglist,
    unsigned long node_num,
    unsigned long zones_count
)
{
    /// 计算每个pglist储存内存区域结构体的数组所需要的内存空间。
    /// 每个内存区域结构体都会分别在每个pglist中被引用，但顺序不同。
    size_t sizeof_zones_each_pglist = sizeof(struct zone_struct*) * zones_count;
    struct zone_struct **tmp_zones = early_alloc(sizeof_zones_each_pglist);
    if (tmp_zones == NULL) {
        /// 单个pglist的内存区域结构体数组分配失败。
        return -1;
    }
    memset(tmp_zones, 0, sizeof_zones_each_pglist);
    pglist->zones = tmp_zones;
    pglist->zones_count = zones_count;
    pglist->node_num = node_num;
    pglist->start_pfn = get_start_pfn();
    pglist->spanned_pages = get_spanned_pages_num();
    init_spinlock_t(&pglist->lock);

    return 0;
}

/**
 * @brief 初始化各个 pglist 结构体的 zones 成员。
 * 
 * 所有内存区域结构体都初始化完成后，初始化 pglist 结构体的 zones 成员。
 * 
 * 排序的规则是：
 * -    在全局 zones 数组中越靠前则在当前 zones 数组中越靠前。
 * -    属于当前 pglist 的内存区域结构体最靠前。
 * -    不单独属于某个 pglist 的内存区域最靠后。
 */
static inline int __init_sort_pglist_zones_array(
    unsigned long valid_zones
)
{
    zones_count = valid_zones;
    for (int node = 0; node < pglists_count; node++) {
        struct pglist_struct *pglist = pglists + node;
        int pglist_zones_count = 0;
        unsigned long pglist_owned_pages = 0;
        unsigned long pglist_valid_pages = 0;
        /// 属于当前 pglist 的内存区域结构体
        for (int i = 0; i < valid_zones; i++) {
            if (zones[i].global_info == pglist) {
                pglist->zones[pglist_zones_count] = zones + i;
                pglist_zones_count++;
                pglist_owned_pages += zones[i].owned_pages;
                pglist_valid_pages += zones[i].owned_pages;
            }
        }
        /// 不属于当前 pglist 的内存区域结构体
        for (int i = 0; i < valid_zones; i++) {
            if (zones[i].global_info != pglist &&
                zones[i].global_info != NULL
            ) {
                pglist->zones[pglist_zones_count] = zones + i;
                pglist_zones_count++;
                pglist_valid_pages += zones[i].owned_pages;
            }
        }
        /// 不单独属于某个 pglist 的内存区域
        for (int i = 0; i < valid_zones; i++) {
            if (zones[i].global_info == NULL) {
                pglist->zones[pglist_zones_count] = zones + i;
                pglist_zones_count++;
                pglist_valid_pages += zones[i].owned_pages;
            }
        }
        pglists->zones_count = pglist_zones_count;
        pglist->owned_pages = pglist_owned_pages;
        pglist->valid_pages = pglist_valid_pages;
    }

    return 0;
}

int pglist_init(unsigned long nodes)
{
    if (nodes == 0) {
        /// 至少要有一个NUMA节点。
        return -1;
    }
    unsigned long areas_count = get_physical_areas_count();
    if (areas_count == 0) {
        /// 至少需要有一个内存区域被添加。
        return -1;
    }
    
    /// 计算总共需要多少个内存区域结构体。因为目前没有实现NUMA，所以需要的数量与 
    /// areas_count 一致。
    unsigned long zones_needed = areas_count;
    /// 事先分配整个 zones 数组需要的空间。
    size_t sizeof_zones = sizeof(struct zone_struct) * areas_count;
    struct zone_struct *tmp_zones = early_alloc(sizeof_zones);
    if (tmp_zones == NULL) {
        return -1;
    }
    zones = tmp_zones;
    
    /// 开始初始化每个pglist
    for (int i = 0; i < nodes; i++) {
        if (__init_pglist_struct(pglists + i, i, zones_needed) != 0) {
            /// 某个pglist结构体初始化失败。
            return -1;
        }
    }

    pglists_count = nodes;
    zones_count = zones_needed;
    return 0;
}

unsigned long get_pglists_count(void)
{
    return pglists_count;
}

struct pglist_struct *get_node_pglist(unsigned long node_num)
{
    if (pglists == NULL || node_num >= pglists_count) {
        return NULL;
    }
    return pglists + node_num;
}

struct pglist_struct *get_current_pglist(void)
{
    /// 当前没有实现NUMA，所以总是返回同一个 pglist 结构体。
    return get_node_pglist(0);
}


/**
 * @brief 判断某个内存区域是否需要为不同的NUMA节点进行分割。
 * 
 * 根据内存区域的大小，判断某个内存区域是否需要分割为多个小区域，让每个的NUMA节点都有
 * 该区域内属于自己的一部分。
 * 
 * 由于当前没有实现NUMA，所以总是返回false。
 * 
 * @retval true 需要分割。
 * @retval false 不需要分割。
 */
static inline int __physical_area_need_splition(
    unsigned long area_start_pfn,
    unsigned long area_end_pfn
)
{
    /// 由于当前没有实现NUMA，所以总是返回false。
    return 0;
}

/**
 * @brief 初始化单个 free_area 结构体。
 */
static inline void __init_free_area_struct(struct free_area *free_area)
{
    atomic_init(&free_area->free_count, 0);
    list_head_init(&free_area->free_list);
}

/**
 * @brief 初始化单个 lruvec 结构体。
 */
static inline void __init_lruvec_struct(struct lruvec *lruvec)
{
    list_head_init(&lruvec->lru_list);
}

/**
 * @brief 初始化单个 zone 结构体。
 */
static inline int __init_zone_struct(
    struct zone_struct *zone,
    struct pglist_struct *pglist,
    unsigned long zone_start_pfn,
    unsigned long zone_spanned_pages
)
{
    memset(zone, 0, sizeof(struct zone_struct));
    __init_lruvec_struct(&zone->lruvec);
    for (int i = 0; i < PAGE_MAX_ORDER; i++) {
        __init_free_area_struct(zone->free_area + i);
    }
    zone->owned_pages = zone_spanned_pages;
    zone->zone_start_pfn = zone_start_pfn;
    zone->global_info = pglist;
    init_spinlock_t(&zone->lock);

    struct page_struct *tmp_pages = \
        physical_to_page(zone_start_pfn << PAGE_SHIFT);
    if (tmp_pages == NULL) {
        return -1;
    }
    /// 使范围内所有的页结构体指向内存区域结构体。
    for (unsigned long i = 0; i < zone_spanned_pages; i++) {
        if (!is_page_valid(tmp_pages + i)) {
            return -1;
        }
        if (tmp_pages[i].zone != NULL) {
            /// 出现内存区域划分的重叠。
            return -1;
        }
        tmp_pages[i].zone = zone;
    }

    return 0;
}

int mem_zones_init(void)
{
    struct pglist_struct *pglist = get_current_pglist();
    if (pglist == NULL) {
        /// pglist 必须被事先初始化。
        return -1;
    }
    if (zones_count == 0 || zones == NULL) {
        /// 内存区域结构体的空间必须事先由 pglist_init 全部得到分配。
        return -1;
    }

    unsigned long page_mask = (1UL << PAGE_SHIFT) - 1;

    /// 用于在遍历内存区域时记录总共有效的内存页数量。
    unsigned long valid_pages_count = 0;
    /// 第一个有效的内存区域的第一个完整的内存页页帧号。
    unsigned long start_pfn = get_start_pfn();
    /// 所有有效的内存区域从最开始到最末尾横跨的完整的内存页数量。
    unsigned long spanned_pages = get_spanned_pages_num();
    /// 所有有效的内存区域之后的第一个有效范围外的内存页页帧号。
    unsigned long end_pfn = start_pfn + spanned_pages;

    /// 开始遍历各个内存区域，并初始化各个内存区域结构体。
    /// 这里假定没有不同的内存区域相互重叠，但是处理过程中并不会检查这一点。
    unsigned long current_zone = 0;
    unsigned long pos;
    unsigned long physical_start;
    unsigned long area_length;
    for_each_physical_area(pos, physical_start, area_length) {
        /// 计算当前内存区域内第一个内存页的页帧号。
        unsigned long area_pfn_start =  \
            (physical_start + page_mask) >> PAGE_SHIFT;
        /// 计算当前内存区域内最后一个内存页的页帧号。
        unsigned long area_pfn_end =    \
            (physical_start + area_length) >> PAGE_SHIFT;
        /// 当前内存区域包括的完整的内存页数量。
        unsigned long area_pages_num = area_pfn_end - area_pfn_start;
        if (area_pfn_start >= area_pfn_end) {
            /// 遍历到的内存区域起始地址冲突，视为初始化失败。
            return -1;;
        }
        if (area_pfn_start < start_pfn || area_pfn_end > end_pfn) {
            /// 对于明显越界的内存区域，视为初始化失败。
            return -1;
        }

        /// 根据是否需要分割内存区域，采用不同的方式进行初始化。
        if (__physical_area_need_splition(area_pfn_start, area_pfn_end)) {
            unsigned long __split_size = area_pages_num + pglists_count - 1;
            __split_size /= pglists_count;
            unsigned long zone_pfn_start = area_pfn_start;
            unsigned long zone_pfn_end = area_pfn_start;
            unsigned long zone_spanned_pages = __split_size;
            for (int i = 0; i < pglists_count; i++) {
                zone_pfn_start = zone_pfn_end;
                zone_pfn_end += zone_spanned_pages;
                if (zone_pfn_end > area_pfn_end) {
                    zone_spanned_pages = area_pfn_end - zone_pfn_start;
                }
                int ret = __init_zone_struct(
                    zones + current_zone, 
                    pglists + i, 
                    zone_pfn_start, 
                    zone_spanned_pages
                );
                if (ret != 0) {
                    return -1;
                } 
                current_zone++;
            }
        } else {
            int ret = __init_zone_struct(
                zones + current_zone, 
                NULL, 
                area_pfn_start, 
                area_pages_num
            );
            if (ret != 0) {
                return -1;
            } 
            current_zone++;
        }

        valid_pages_count += area_pages_num;
    }

    /// 所有内存区域结构体都初始化完成后，初始化 pglist 结构体的 zones 成员。
    return __init_sort_pglist_zones_array(current_zone);
}

unsigned long get_global_zones_count(void)
{
    if (zones == NULL) {
        return 0;
    }
    return zones_count;
}

unsigned long get_pglist_zones_count(struct pglist_struct *pglist)
{
    if (pglist == NULL) {
        return 0;
    }
    return pglist->zones_count;
}

struct zone_struct *get_global_zone(unsigned long num)
{
    if (num >= get_global_zones_count()) {
        return NULL;
    }
    return zones + num;
}

struct zone_struct *get_pglist_zone(
    struct pglist_struct *pglist, 
    unsigned long num
)
{
    unsigned long paglist_zones_count = get_pglist_zones_count(pglist);
    if (num >= paglist_zones_count) {
        return NULL;
    }
    return pglist->zones[num];
}
